import sys
from PyQt5.QtCore import Qt, QTimer
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QVBoxLayout, QWidget, QPushButton, QHBoxLayout, \
    QMessageBox, QFileDialog, QProgressBar
from PyQt5.QtGui import QImage, QPixmap, QIcon
import cv2
from ultralytics import YOLO


class Worker:
    def __init__(self):
        self.model = None

    def load_model(self):
        model_path, _ = QFileDialog.getOpenFileName(None, "选择模型文件", "", "模型文件 (*.pt)")
        if model_path:
            self.model = YOLO(model_path)
            return self.model is not None
        return False

    def detect_image(self, image):
        results = self.model.predict(image)
        return results


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Altria的农作物病害检测系统")
        # self.setWindowIcon(QIcon("icon.png"))
        self.setGeometry(300, 150, 800, 400)

        # 创建两个 QLabel 分别显示左右图像
        self.label1 = QLabel()
        self.label1.setAlignment(Qt.AlignCenter)
        self.label1.setMinimumSize(580, 450)  # 设置大小
        self.label1.setStyleSheet('border:3px solid #6950a1; background-color: black;')  # 添加边框并设置背景颜色为黑色

        self.label2 = QLabel()
        self.label2.setAlignment(Qt.AlignCenter)
        self.label2.setMinimumSize(580, 450)  # 设置大小
        self.label2.setStyleSheet('border:3px solid #6950a1; background-color: black;')  # 添加边框并设置背景颜色为黑色

        # 水平布局，用于放置左右两个 QLabel
        layout = QVBoxLayout()
        # layout.addWidget(self.label1)
        hbox_video = QHBoxLayout()
        hbox_video.addWidget(self.label1)  # 左侧显示原始图像
        hbox_video.addWidget(self.label2)  # 右侧显示检测后的图像
        layout.addLayout(hbox_video)
        self.worker = Worker()
        # 创建按钮布局
        hbox_buttons = QHBoxLayout()
        # 添加模型选择按钮
        self.load_model_button = QPushButton("📁模型选择")
        self.load_model_button.clicked.connect(self.load_model)
        self.load_model_button.setFixedSize(120, 30)
        hbox_buttons.addWidget(self.load_model_button)

        # 添加图片检测按钮
        self.image_detect_button = QPushButton("💾图片检测")
        self.image_detect_button.clicked.connect(self.detect_image)
        self.image_detect_button.setEnabled(False)
        self.image_detect_button.setFixedSize(120, 30)
        hbox_buttons.addWidget(self.image_detect_button)

        # 添加显示检测物体按钮
        self.display_objects_button = QPushButton("🔍显示检测物体")
        self.display_objects_button.clicked.connect(self.show_detected_objects)
        self.display_objects_button.setEnabled(False)
        self.display_objects_button.setFixedSize(120, 30)
        hbox_buttons.addWidget(self.display_objects_button)

        # 添加视频流检测按钮
        self.video_detect_button = QPushButton("📹视频流检测")
        self.video_detect_button.clicked.connect(self.toggle_video_detection)
        self.video_detect_button.setEnabled(False)
        self.video_detect_button.setFixedSize(120, 30)
        hbox_buttons.addWidget(self.video_detect_button)

        # 添加暂停按钮
        self.pause_button = QPushButton("⏸暂停")
        self.pause_button.clicked.connect(self.pause_video)
        self.pause_button.setEnabled(False)
        self.pause_button.setFixedSize(120, 30)
        hbox_buttons.addWidget(self.pause_button)

        # 添加继续播放按钮
        self.resume_button = QPushButton("▶️继续")
        self.resume_button.clicked.connect(self.resume_video)
        self.resume_button.setEnabled(False)
        self.resume_button.setFixedSize(120, 30)
        hbox_buttons.addWidget(self.resume_button)

        # 添加进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setFixedSize(580, 20)
        layout.addWidget(self.progress_bar)

        # 添加退出按钮
        self.exit_button = QPushButton("❌退出")
        self.exit_button.clicked.connect(self.exit_application)
        self.exit_button.setFixedSize(120, 30)
        hbox_buttons.addWidget(self.exit_button)

        layout.addLayout(hbox_buttons)
        central_widget = QWidget()
        central_widget.setLayout(layout)
        self.setCentralWidget(central_widget)

        self.current_results = None

        self.timer = QTimer()
        self.timer.timeout.connect(self.process_video_frame)
        self.cap = None

    def detect_image(self):
        image_path, _ = QFileDialog.getOpenFileName(None, "选择图片文件", "", "图片文件 (*.jpg *.jpeg *.png)")
        if image_path:
            image = cv2.imread(image_path)
            if image is not None:
                self.current_results = self.worker.detect_image(image)
                if self.current_results:
                    annotated_image = self.current_results[0].plot()
                    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 转换为 RGB
                    height1, width1, channel1 = image_rgb.shape
                    bytesPerLine1 = 3 * width1
                    qimage1 = QImage(image_rgb.data, width1, height1, bytesPerLine1, QImage.Format_RGB888)
                    pixmap1 = QPixmap.fromImage(qimage1)
                    self.label1.setPixmap(pixmap1.scaled(self.label1.size(), Qt.KeepAspectRatio))

                    annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)  # 转换为 RGB
                    height2, width2, channel2 = annotated_image.shape
                    bytesPerLine2 = 3 * width2
                    qimage2 = QImage(annotated_image.data, width2, height2, bytesPerLine2, QImage.Format_RGB888)
                    pixmap2 = QPixmap.fromImage(qimage2)
                    self.label2.setPixmap(pixmap2.scaled(self.label2.size(), Qt.KeepAspectRatio))

    def show_detected_objects(self):
        if self.current_results:
            det_info = self.current_results[0].boxes.cls
            object_count = len(det_info)
            object_info = f"识别到的物体总个数：{object_count}\n"
            object_dict = {}
            class_names_dict = self.current_results[0].names
            for class_id in det_info:
                class_name = class_names_dict[int(class_id)]
                if class_name in object_dict:
                    object_dict[class_name] += 1
                else:
                    object_dict[class_name] = 1
            sorted_objects = sorted(object_dict.items(), key=lambda x: x[1], reverse=True)
            for obj_name, obj_count in sorted_objects:
                object_info += f"{obj_name}: {obj_count}\n"
            self.show_message_box("识别结果", object_info)
        else:
            self.show_message_box("识别结果", "未检测到物体")

    def show_message_box(self, title, message):
        msg_box = QMessageBox(self)
        msg_box.setWindowTitle(title)
        msg_box.setText(message)
        msg_box.exec_()

    def load_model(self):
        if self.worker.load_model():
            self.image_detect_button.setEnabled(True)
            self.video_detect_button.setEnabled(True)
            self.display_objects_button.setEnabled(True)

    def exit_application(self):
        # 终止程序运行
        sys.exit()

    def toggle_video_detection(self):
        if self.timer.isActive():
            self.timer.stop()
            if self.cap:
                self.cap.release()
            self.video_detect_button.setText("📹视频流检测")
            self.pause_button.setEnabled(False)
            self.resume_button.setEnabled(False)
        else:
            video_path, _ = QFileDialog.getOpenFileName(None, "选择视频文件", "", "视频文件 (*.mp4 *.avi)")
            if video_path:
                self.cap = cv2.VideoCapture(video_path)
                if self.cap.isOpened():
                    self.timer.start(30)  # 每30毫秒处理一帧
                    self.video_detect_button.setText("⏹停止检测")
                    self.pause_button.setEnabled(True)
                    self.resume_button.setEnabled(False)
                    self.progress_bar.setMaximum(int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT)))

    def process_video_frame(self):
        ret, frame = self.cap.read()
        if ret:
            self.current_results = self.worker.detect_image(frame)
            if self.current_results:
                annotated_frame = self.current_results[0].plot()
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                height1, width1, channel1 = frame_rgb.shape
                bytesPerLine1 = 3 * width1
                qimage1 = QImage(frame_rgb.data, width1, height1, bytesPerLine1, QImage.Format_RGB888)
                pixmap1 = QPixmap.fromImage(qimage1)
                self.label1.setPixmap(pixmap1.scaled(self.label1.size(), Qt.KeepAspectRatio))

                annotated_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
                height2, width2, channel2 = annotated_frame.shape
                bytesPerLine2 = 3 * width2
                qimage2 = QImage(annotated_frame.data, width2, height2, bytesPerLine2, QImage.Format_RGB888)
                pixmap2 = QPixmap.fromImage(qimage2)
                self.label2.setPixmap(pixmap2.scaled(self.label2.size(), Qt.KeepAspectRatio))

            # 更新进度条
            current_frame = int(self.cap.get(cv2.CAP_PROP_POS_FRAMES))
            self.progress_bar.setValue(current_frame)
        else:
            self.timer.stop()
            self.cap.release()
            self.video_detect_button.setText("📹视频流检测")
            self.pause_button.setEnabled(False)
            self.resume_button.setEnabled(False)

    def pause_video(self):
        if self.timer.isActive():
            self.timer.stop()
            self.pause_button.setEnabled(False)
            self.resume_button.setEnabled(True)

    def resume_video(self):
        if not self.timer.isActive() and self.cap:
            self.timer.start(30)
            self.pause_button.setEnabled(True)
            self.resume_button.setEnabled(False)


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())
